package org.zry.chat2ai.controller;

import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.util.CollectionUtils;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import org.zry.chat2ai.entity.ChatHistory;
import reactor.core.Disposable;

import javax.validation.Valid;
import javax.validation.constraints.NotNull;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@RestController
public class ChatController {
    private final OpenAiChatModel chatModel;
    @Autowired
    private MongoTemplate mongoTemplate;

    @Autowired
    public ChatController(OpenAiChatModel chatModel) {
        this.chatModel = chatModel;
    }

    @GetMapping("/ai/generate")
    public Map generate(@RequestParam(value = "message", defaultValue = "你是谁？") String message) {
        return Map.of("generation", this.chatModel.call(message));
    }

    @GetMapping("/ai/generateStream")
    public SseEmitter streamChat(@RequestParam @Valid @NotNull String userId, @RequestParam @Validated @NotNull String sessionId, @RequestParam @Valid @NotNull String message) {
        // 创建 SSE 发射器，设置超时时间（例如 1 分钟）
        SseEmitter emitter = new SseEmitter(6000000L);
        // 查找或创建会话记录
        Query query = new Query(Criteria.where("userId").is(userId).and("sessionId").is(sessionId));
        ChatHistory chatHistory = mongoTemplate.findOne(query, ChatHistory.class, "chat_history");
        if (chatHistory == null) {
            chatHistory = new ChatHistory(userId, sessionId, new ArrayList<>());
        }
        // 创建 Prompt 对象
        chatHistory.getUserMessages().add(message);
        StringBuilder sb = new StringBuilder();
        Prompt prompt = new Prompt(chatHistory.getMessagesByString());
        // 订阅流式响应
        ChatHistory finalChatHistory = chatHistory;
        Disposable subscribe = chatModel.stream(prompt).subscribe(response -> {
                    try {
                        AssistantMessage output = response.getResult().getOutput();
                        String content = response.getResult().getOutput().getText();
                        System.out.print(content);
                        sb.append(content);
                        // 发送 SSE 事件
                        emitter.send(SseEmitter.event()
                                .data(content)
                                .id(String.valueOf(System.currentTimeMillis()))
                                .build());
                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                },
                error -> { // 异常处理
                    emitter.completeWithError(error);
                },
                () -> { // 完成处理
                    String assistantResponse = sb.toString();
                    if (!assistantResponse.isEmpty()) {
                        assistantResponse = assistantResponse.replaceAll("<think>[\\s\\S]*?</think>", "");
                        System.out.println("assistantResponse=" + assistantResponse);
                        finalChatHistory.getUserMessages().add(assistantResponse);
                        mongoTemplate.save(finalChatHistory);
                    }
                    emitter.complete();
                }
        );
        // 处理客户端断开连接
        emitter.onCompletion(() -> {
            subscribe.dispose();
            // 可在此处释放资源
            System.out.println("SSE connection completed");
        });
        emitter.onTimeout(() -> {
            emitter.complete();
            System.out.println("SSE connection timed out");
        });
        return emitter;
    }

    @GetMapping("/ai/chatHistoryList")
    public List<ChatHistory> chatHistory(@RequestParam @Valid String userId) {
        Query query = new Query(Criteria.where("userId").is(userId));
        List<ChatHistory> chatHistorys = mongoTemplate.find(query, ChatHistory.class);
        if (CollectionUtils.isEmpty(chatHistorys)) {
            return new ArrayList<>();
        }
        return chatHistorys;
    }
}